//
// Created by DonAzufre on 2024/4/1.
//
#pragma once

#include <array>
#include <string>
#include "IR.h"
#include "IRBuilder.h"


const std::array<std::string, 13> libFuncs = {
    "getint",   "getch",    "getfloat",  "getarray",           "getfarray",       "putint",        "putch",
    "putfloat", "putarray", "putfarray", "llvm.memset.p0.i32", "_sysy_starttime", "_sysy_stoptime"};

class Pass
{
public:
    std::string passName;

    virtual sysy::Module *run(sysy::Module *) = 0;

    virtual ~Pass() = default;
};

class PreProcess : public Pass
{
public:
    PreProcess() { passName = "PreProcess"; }

    sysy::Module *run(sysy::Module *module) override
    {
        // module = AddRetVoid(module);
        module = RemoveBBlocks(module);
        module = MoveAllocaToEntry(module);
        module = RenameReg(module, false);
        return module;
    }

    // static sysy::Module *AddRetVoid(sysy::Module *module)
    // {
    //     sysy::IRBuilder builder;
    //     for (const auto &[name, func] : module->getFunctions())
    //     {
    //         if (func->getReturnType()->isVoid())
    //         {
    //             auto last = func->getExitBlock();
    //             if (last->getInstructions().empty() || !last->getLastInst()->isRet())
    //             {
    //                 builder.setPosition(last, last->end());
    //                 builder.createReturnInst(nullptr);
    //             }
    //         }
    //     }
    //     return module;
    // }

    static sysy::Module *RemoveBBlocks(sysy::Module *module)
    {
        // 删除空基本块
        // 删除没有前驱的基本块
        auto &funcs = module->getFunctions();

        std::cerr << "Removed bloks:" << std::endl;

        for (const auto &[name, func] : funcs)
        {
            if (std::find(libFuncs.begin(), libFuncs.end(), name) != libFuncs.end())
                continue;
            std::vector<sysy::BasicBlock *> emptyBBs;
            for (auto &bb : func->getBBlocks())
            {
                if (bb->isEmptyBBlock())
                {
                    // 空基本块前驱的最后一条指令一定是跳转, 目标是该空块
                    // 调整到跳转到空块后继

                    // 空块一定只有一个后继? 大概
                    assert(bb->getSuccessors().size() == 1);
                    auto succ = bb->getSuccessors()[0];

                    for (auto pre : bb->getPredecessors())
                    {
                        auto br = pre->getLastInst();
                        assert(br->isBranch());
                        for (auto &op : br->getOperands())
                        {
                            if (op.getValue() == static_cast<sysy::Value *>(bb.get()))
                            {
                                br->replaceOperand(op.getIndex(), succ);
                                auto _ =
                                    std::remove(pre->getSuccessors().begin(), pre->getSuccessors().end(), bb.get());
                                pre->getSuccessors().emplace_back(succ);
                                succ->getPredecessors().emplace_back(pre);
                                break;
                            }
                        }
                    }
                    emptyBBs.push_back(bb.get());

                    // if (bb->getPredecessors().empty())
                    // {
                    //     emptyBBs.push_back(bb.get());
                    // }
                    // else
                    // {
                    //     for (auto pre : bb->getPredecessors())
                    //     {
                    //         auto _ = std::remove(pre->getSuccessors().begin(), pre->getSuccessors().end(), bb.get());
                    //     }
                    //     emptyBBs.push_back(bb.get());
                    // }
                }
                else if (bb->getPredecessors().empty())
                {
                    if (bb.get() == func->getEntryBlock())
                        continue;
                    for (auto suc : bb->getSuccessors())
                    {
                        auto _ = std::remove(suc->getPredecessors().begin(), suc->getPredecessors().end(), bb.get());
                    }
                    emptyBBs.push_back(bb.get());
                }
            }

            for (auto eBb : emptyBBs)
            {
                std::cerr << eBb->getName() << std::endl;
                func->removeBasicBlock(eBb);
            }
        }
        return module;
    }


    static sysy::Module *RenameReg(sysy::Module *pm, bool meaningfulBBName = true)
    {
        for (auto &[name, func] : pm->getFunctions())
        {
            int reg_num = 0;
            auto new_reg_name = [&reg_num]() { return std::string("%") + std::to_string(reg_num++); };

            for (auto &arg : func->getArguments())
            {
                arg.setName(new_reg_name());
            }

            for (auto &blk : func->getBasicBlocks())
            {
                // inst和basicblock共享计数
                if (!meaningfulBBName)
                    blk->setName(std::to_string(reg_num));
                reg_num++;

                for (auto &inst : blk->getInstructions())
                {
                    if (inst->noRegName())
                        continue;
                    inst->setName(new_reg_name());
                }
            }
        }
        return pm;
    }

    static sysy::Module *MoveAllocaToEntry(sysy::Module *pm)
    {
        for (auto &[name, func] : pm->getFunctions())
        {
            if (find(libFuncs.begin(), libFuncs.end(), name) != libFuncs.end())
                continue;
            auto entry = func->getEntryBlock();
            std::list<std::unique_ptr<sysy::Instruction>> allocas;
            for (auto &blk : func->getBasicBlocks())
            {
                for (auto &inst : blk->getInstructions())
                {
                    if (inst->isAlloca())
                    {
                        allocas.emplace_back(inst.release());
                    }
                }
                blk->getInstructions().remove_if([&](std::unique_ptr<sysy::Instruction> &up) { return up == nullptr; });
            }
            entry->getInstructions().splice(entry->getInstructions().begin(), allocas);
        }
        return pm;
    }
};

class Mem2Reg : public Pass
{
protected:
    /**
     * parent(w): The vertex which is the parent of vertex w in the spanning tree generated by the search.
     * semi(w):A number defmed as follows:
        (i) Before vertex w is numbered, semi(v) = O.
        (ii) After w is numbered but before its semidominator is computed, semi(w) is the number of w.
        (iii) After the semidominator of w is computed, semi(w) is the number of the semidominator of w.
     * vertex (i):The vertex whose number is i.
     * bucket(w): A set of vertices whose semidominator is w.
     * idom(w): A vertex defined as follows:
        (i) After step 3, if the semidominator of w is its immediate dominator, then dom(w) is
        the immediate dominator of w. Otherwise dom(w) is a vertex v whose number is
        smaller than w and whose immediate dominator is also w's immediate dominator.
        (ii) After step 4, dom(w) is the immediate dominator of w.
     *
     * **/
    std::map<sysy::BasicBlock *, sysy::BasicBlock *> parentNode;
    std::map<sysy::BasicBlock *, int> semi;
    std::map<int, sysy::BasicBlock *> vertex;
    std::map<sysy::BasicBlock *, int> dfn;
    std::map<sysy::BasicBlock *, std::list<sysy::BasicBlock *>> bucket;
    std::map<sysy::BasicBlock *, sysy::BasicBlock *> idom;
    std::map<sysy::BasicBlock *, int> ancestor;
    std::map<sysy::BasicBlock *, sysy::BasicBlock *> label;
    std::map<sysy::BasicBlock *, int> size;
    std::map<sysy::BasicBlock *, int> child;
    std::map<sysy::BasicBlock *, std::list<sysy::BasicBlock *>> domTree;
    int n;

public:
    Mem2Reg() { passName = "Mem2Reg"; }

    sysy::Module *run(sysy::Module *module) override
    {
        module = Mem2Reg_Phi(module);
        module = PreProcess::RenameReg(module, false);
        return module;
    }

    void DFS(sysy::BasicBlock *root)
    {
        n += 1;
        semi[root] = n;
        vertex[n] = root;
        dfn[root] = n;
        ancestor[root] = 0;
        label[root] = root;
        size[root] = 1;
        child[root] = 0;
        for (auto &bb : root->getSuccessors())
        {
            if (semi[bb] == 0)
            {
                parentNode[bb] = root;
                DFS(bb);
            }
        }
    }

    void COMPRESS(sysy::BasicBlock *v)
    {
        if (ancestor[vertex[ancestor[v]]] != 0)
        {
            COMPRESS(vertex[ancestor[v]]);
            if (semi[label[vertex[ancestor[v]]]] < semi[label[v]])
            {
                label[v] = label[vertex[ancestor[v]]];
            }
            ancestor[v] = ancestor[vertex[ancestor[v]]];
        }
    }

    sysy::BasicBlock *EVAL(sysy::BasicBlock *v)
    {
        sysy::BasicBlock *EVAL_ANS;
        if (ancestor[v] == 0)
        {
            EVAL_ANS = label[v];
        }
        else
        {
            COMPRESS(v);
            if (semi[label[vertex[ancestor[v]]]] >= semi[label[v]])
            {
                EVAL_ANS = label[v];
            }
            else
            {
                EVAL_ANS = label[vertex[ancestor[v]]];
            }
        }
        return EVAL_ANS;
    }

    void LINK(sysy::BasicBlock *v, sysy::BasicBlock *w)
    {
        auto s = w;
        while (semi[label[w]] < semi[label[vertex[child[s]]]])
        {
            if (size[s] + size[vertex[child[vertex[child[s]]]]] >= 2 * (size[vertex[child[s]]]))
            {
                parentNode[vertex[child[s]]] = s;
                child[s] = child[vertex[child[s]]];
            }
            else
            {
                size[vertex[child[s]]] = size[s];
                parentNode[s] = vertex[child[s]];
                s = parentNode[s];
            }
        }
        label[s] = label[w];
        size[v] = size[v] + size[w];
        if (size[v] < 2 * size[w])
        {
            auto tmp = s;
            s = vertex[child[v]];
            child[v] = dfn[s];
        }
        while (dfn[s] != 0)
        {
            parentNode[s] = v;
            s = vertex[child[s]];
        }
    }

    void LT(sysy::Function *func)
    {
        // step1:
        n = 0;
        for (auto &bb : func->getBasicBlocks())
        {
            semi[bb.get()] = 0;
        }
        DFS(func->getBasicBlocks().begin()->get());

        // step2 and step3:
        for (int i = n; i >= 2; i--)
        {
            auto w = vertex[i];
            for (auto v : w->getPredecessors())
            {
                auto u = EVAL(v);
                if (semi[u] < semi[w])
                {
                    semi[w] = semi[u];
                }
            }
            bucket[vertex[semi[w]]].emplace_back(w);
            LINK(parentNode[w], w);
            for (auto v : bucket[parentNode[w]])
            {
                //                bucket[parentNode[w]].remove(v);
                auto u = EVAL(v);
                if (semi[u] < semi[v])
                {
                    idom[v] = u;
                }
                else
                {
                    idom[v] = parentNode[w];
                }
            }
        }

        // step4:
        for (int i = 2; i <= n; i++)
        {
            auto w = vertex[i];
            if (idom[w] != vertex[semi[w]])
            {
                idom[w] = idom[idom[w]];
            }
        }
        idom[func->getBasicBlocks().begin()->get()] = nullptr;

        for (auto idomMap : idom)
        {
            domTree[idomMap.second].emplace_back(idomMap.first);
        }
    }

    //    /**计算b块直接支配的节点**/
    //    static std::list<sysy::BasicBlock*> findDominate(sysy::BasicBlock *root, sysy::BasicBlock *targetNode){
    //        return {};
    //    }
    //
    //    /**计算b块被直接支配的节点**/
    //    static std::list<sysy::BasicBlock*> findIDominated(sysy::BasicBlock *root, sysy::BasicBlock *targetNode){
    //        return {};
    //    }

    /**单独块对%v做如下优化：
     * 线性扫描基本块S，将每一个load替换为前一个store的值，删除该load
     * 若在load之前没有发生对%v的store，则不删除该load
     * 若一个store后面没有任何对%v的load，则删除该store**/
    static std::map<sysy::Value *, sysy::Value *>
    constOneBlock(std::map<sysy::Value *, sysy::Value *> rm, sysy::BasicBlock *bb, std::list<sysy::Value *> allocaRegs)
    {
        std::unique_ptr<sysy::Instruction> store_end = nullptr;
        std::map<int, sysy::Value *> regUpdataMaps;
        std::map<int, int> idexUpdataMaps;
        std::map<int, sysy::Instruction *> insUpdataMaps;
        int i = 0;
        for (auto &ins : bb->getInstructions())
        {
            for (auto allocaReg : allocaRegs)
            {
                if (ins->isLoad() && ins->getOperand(0) == allocaReg)
                {
                    if (rm.find(allocaReg) != rm.end())
                    {
                        if (rm[allocaReg] != allocaReg)
                        {
                            rm[ins.get()] = rm[allocaReg];
                            auto remove_ins = ins.release();
                        }
                    }
                }
                else if (ins->isStore() && ins->getOperand(1) == allocaReg)
                {
                    rm[allocaReg] = ins->getOperand(0);
                    store_end = std::unique_ptr<sysy::Instruction>(ins.release());
                }
                else if (ins->isAlloca())
                {
                    //
                }
                else
                {
                    int j = 0;
                    for (auto ope : ins->getOperands())
                    {
                        if (rm.find(ope.getValue()) != rm.end())
                        {
                            regUpdataMaps[i] = rm[ope.getValue()];
                            idexUpdataMaps[i] = j;
                            insUpdataMaps[i] = ins.get();
                            if (std::find(allocaRegs.begin(), allocaRegs.end(), ope.getValue()) == allocaRegs.end())
                            {
                                rm.erase(ope.getValue());
                            }
                            i++;
                        }
                        j++;
                    }
                }
            }
        }
        for (int j = 0; j < i; j++)
        {
            insUpdataMaps[j]->replaceOperand(idexUpdataMaps[j], regUpdataMaps[j]);
        }
        if (store_end)
        {
            bb->getInstructions().emplace((--bb->getInstructions().end()),
                                          std::unique_ptr<sysy::Instruction>(store_end.release()));
        }
        bb->getInstructions().remove_if([&](std::unique_ptr<sysy::Instruction> &up) { return up == nullptr; });
        return rm;
    }

    // 因为reg_Value里的Alloca值只在一个块中有store，所以里面Alloca不会变化了
    std::map<sysy::Value *, sysy::Value *> oneDefineBlock(std::map<sysy::Value *, sysy::Value *> rm,
                                                          sysy::BasicBlock *bb, std::list<sysy::Value *> allocaRegs)
    {
        std::map<sysy::Value *, sysy::Value *> reg_Value;
        for (auto bbSucc : domTree[bb])
        {
            reg_Value = constOneBlock(rm, bbSucc, allocaRegs);
            reg_Value = oneDefineBlock(reg_Value, bbSucc, allocaRegs);
        }
        return reg_Value;
    }

    sysy::Module *Mem2Reg_Phi(sysy::Module *module)
    {
        auto &funcs = module->getFunctions();
        for (const auto &[name, func] : funcs)
        {
            if (!func->getBasicBlocks().empty() &&
                std::find(libFuncs.begin(), libFuncs.end(), func->getName()) == libFuncs.end())
            {
                LT(func.get());
            }
            int alloca_num = 0;
            // 代码申请全部alloca的信息栈
            static std::map<int, sysy::Instruction *> allocaListMaps;
            // DefiningBlocks：存在对%v发生store的所有基本块
            static std::map<int, std::list<sysy::BasicBlock *>> DefiningBlocksMaps;
            // UsingBlocks：存在对对%v发生load的所有基本块
            static std::map<int, std::list<sysy::BasicBlock *>> UsingBlocksListMaps;
            // GEPBlocksListMaps：存在对对%v发生GEP的所有基本块
            static std::map<int, std::list<sysy::BasicBlock *>> GEPBlocksListMaps;
            for (auto &bb : func->getBBlocks())
            {
                for (auto &ins : bb->getInstructions())
                {
                    /**对于指令%v = alloca: type，找到以下信息
                     * allocaListMaps：代码申请全部alloca的信息栈
                     * DefiningBlocksMaps：存在对%v发生store的所有基本块
                     * UsingBlocksListMaps：存在对对%v发生load的所有基本块
                     * GEPBlocksListMaps：存在对对%v发生GEP的所有基本块**/
                    if (ins->isAlloca())
                    {
                        allocaListMaps[alloca_num] = ins.release();
                        DefiningBlocksMaps[alloca_num] = {};
                        UsingBlocksListMaps[alloca_num] = {};
                        GEPBlocksListMaps[alloca_num] = {};
                        alloca_num++;
                    }
                    else if (ins->isStore())
                    {
                        int i = 0;
                        while (i < alloca_num)
                        {
                            if (ins->getOperand(1)->getName() == allocaListMaps[i]->getName() &&
                                std::find(DefiningBlocksMaps[i].begin(), DefiningBlocksMaps[i].end(), &*bb) ==
                                    DefiningBlocksMaps[i].end())
                            {
                                DefiningBlocksMaps[i].emplace_back(&*bb);
                                break;
                            }
                            i++;
                        }
                    }
                    else if (ins->isLoad())
                    {
                        int i = 0;
                        while (i < alloca_num)
                        {
                            if (ins->getOperand(0)->getName() == allocaListMaps[i]->getName() &&
                                std::find(UsingBlocksListMaps[i].begin(), UsingBlocksListMaps[i].end(), &*bb) ==
                                    UsingBlocksListMaps[i].end())
                            {
                                UsingBlocksListMaps[i].emplace_back(&*bb);
                                break;
                            }
                            i++;
                        }
                    }
                    else if (ins->isGEP())
                    {
                        int i = 0;
                        while (i < alloca_num)
                        {
                            if (ins->getOperand(0)->getName() == allocaListMaps[i]->getName() &&
                                std::find(GEPBlocksListMaps[i].begin(), GEPBlocksListMaps[i].end(), &*bb) ==
                                    GEPBlocksListMaps[i].end())
                            {
                                GEPBlocksListMaps[i].emplace_back(&*bb);
                                break;
                            }
                            i++;
                        }
                    }
                }
                bb->getInstructions().remove_if([&](std::unique_ptr<sysy::Instruction> &up) { return up == nullptr; });
            }
            auto entry = func->getEntryBlock();
            int i = alloca_num - 1;
            while (i >= 0)
            {
                if (!UsingBlocksListMaps[i].empty() || !GEPBlocksListMaps[i].empty())
                {
                    entry->getInstructions().push_front(std::unique_ptr<sysy::Instruction>(allocaListMaps[i]));
                    if (DefiningBlocksMaps[i].size() == 1 &&
                        std::find(GEPBlocksListMaps[i].begin(), GEPBlocksListMaps[i].end(),
                                  DefiningBlocksMaps[i].front()) == GEPBlocksListMaps[i].end())
                    {
                        /// 如果变量%v的DefiningBlocks中只有一个基本块S
                        std::map<sysy::Value *, sysy::Value *> reg_value;
                        std::list<sysy::Value *> allocaRegs;
                        for (auto bb_Def : DefiningBlocksMaps[i])
                        {
                            allocaRegs.emplace_back(allocaListMaps[i]);
                            reg_value[allocaListMaps[i]] = allocaListMaps[i];
                            reg_value = constOneBlock(reg_value, bb_Def, allocaRegs);
                            reg_value = oneDefineBlock(reg_value, bb_Def, allocaRegs);
                        }
                    }
                    /// 如果变量%v的所有读写均在一个块S内发生
                    for (auto bb_Def : DefiningBlocksMaps[i])
                    {
                        for (auto bb_Use : UsingBlocksListMaps[i])
                        {
                            if (bb_Def == bb_Use &&
                                std::find(GEPBlocksListMaps[i].begin(), GEPBlocksListMaps[i].end(), bb_Def) ==
                                    GEPBlocksListMaps[i].end())
                            {
                                std::map<sysy::Value *, sysy::Value *> reg_value;
                                std::list<sysy::Value *> allocaRegs;
                                allocaRegs.emplace_back(allocaListMaps[i]);
                                reg_value[allocaListMaps[i]] = allocaListMaps[i];
                                reg_value = constOneBlock(reg_value, bb_Def, allocaRegs);
                                break;
                            }
                        }
                    }
                }
                else
                {
                    for (auto bb : DefiningBlocksMaps[i])
                    {
                        for (auto &ins : bb->getInstructions())
                        {
                            if (ins->isStore() && allocaListMaps[i]->getName() == ins->getOperand(1)->getName())
                            {
                                auto remove_ins = ins.release();
                            }
                        }
                        bb->getInstructions().remove_if([&](std::unique_ptr<sysy::Instruction> &up)
                                                        { return up == nullptr; });
                    }
                    allocaListMaps.erase(i);
                    DefiningBlocksMaps.erase(i);
                    UsingBlocksListMaps.erase(i);
                    GEPBlocksListMaps.erase(i);
                }
                i--;
            }
        }
        return module;
    }
};
